/*
 * GeoState.cpp
 *
 *  Created on: 14 Oct 2011
 *      Author: allan
 */

#include "GeoState.h"

// C++ includes
#include <iomanip>

// Boost includes
#include <boost/foreach.hpp>

// GeoReact includes
#include "Geochemistry/ElectrolyteSolution.h"

GeoState::GeoState(const GeoSystem& geosystem) : 
geosystem(geosystem), T(25.0), P(1.0)
{
	n = a = g = r = VectorXd::Constant(geosystem.GetMultiphase().GetSpecies().size(), 1.0E-7);
}

void GeoState::SetTemperature(double T)
{
	this->T = T;
}

void GeoState::SetPressure(double P)
{
	this->P = P;
}

void GeoState::SetComposition(const VectorXd& n)
{
	const Multiphase& multiphase = geosystem.GetMultiphase();
	
	const ReactionSystem& kReactions = geosystem.GetKineticReactions();
	
	MatrixXd uki = kReactions.AssembleStoichiometricMatrix(multiphase.GetSpecies());
	
	this->n = n;
	
	a = multiphase.Activities(T, P, n);
	
	g = multiphase.ActivityCoefficients(T, P, n);
	
	r = uki.transpose() * kReactions.KineticRates(T, P, n, a);
}

void GeoState::SetComposition(const string& species, double val)
{
	n[geosystem.GetMultiphase()[species]] = val;
}

const double GeoState::GetTemperature() const
{
	return T;
}
	
const double GeoState::GetPressure() const
{
	return P;
}

const VectorXd& GeoState::GetComposition() const
{
	return n;
}

const VectorXd GeoState::GetPhaseComposition(const string& phasename) const
{
	unsigned start = 0;
	unsigned size = 0;
	
	BOOST_FOREACH(const Phase& phase, geosystem.GetMultiphase().GetPhases())
	{
		if(phase.GetName() == phasename)
			size = phase.GetSpecies().size();
		else
			start += phase.GetSpecies().size();
	}
	
	return n.segment(start, size);
}

const double GeoState::Get(const string& quantity) const
{
	if(quantity == "IonicStrength" || quantity == "IonicStrength(Effective)")
	{
		const vector<string> aqueousSpecies = geosystem.GetMultiphase().GetPhase("Aqueous").GetSpecies();
		
		ElectrolyteSolution sol(aqueousSpecies);
		
		return sol.EffectiveIonicStrength(GetPhaseComposition("Aqueous"));
	}
	
	if(quantity == "IonicStrength(Stoichiometric)")
	{
		const vector<string> aqueousSpecies = geosystem.GetMultiphase().GetPhase("Aqueous").GetSpecies();
		
		ElectrolyteSolution sol(aqueousSpecies);
		
		return sol.StoichiometricIonicStrength(GetPhaseComposition("Aqueous"));
	}
	
	if(quantity == "pH" || quantity == "Acidity")
	{
		double iH = geosystem.GetMultiphase()["H[+]"];
		
		return -log10(a[iH]);
	}
	
	cerr << "Error: The quantity type " << quantity << " is not valid." << endl; return 0.0;
}

const double GeoState::Get(const string& quantity, const string& species) const
{
	const Index ispecies = geosystem.GetMultiphase()[species];
	
	if(quantity == "Moles" || quantity == "Concentration")
		return n[ispecies];
	
	if(quantity == "Activity")
		return a[ispecies];
	
	if(quantity == "ActivityCoefficient")
		return g[ispecies];
	
	if(quantity == "PartialPressure")
		return P * (n[ispecies] / GetPhaseComposition("Gaseous").sum());
	
	if(quantity == "KineticRate")
		return r[ispecies];
	
	cerr << "Error: The quantity type " << quantity << " is not valid." << endl; return 0.0;
}

void GeoState::OutputMultiphaseState(ostream& out) const
{
	const Multiphase& multiphase = geosystem.GetMultiphase();
	
	out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
	out << ">>>         Multiphase State         <<<" << endl;
	out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
	vector<string> species = multiphase.GetSpecies();
	
	auto compare = [&](const string& s1, const string& s2)
		{ return n[multiphase[s1]] > n[multiphase[s2]]; };
	
	sort(species.begin(), species.end(), compare);
	
	out << setprecision(6);
	
	out << setw(25) << left << "Species";
	out << setw(25) << left << "Moles";
	out << setw(25) << left << "Activity";
	out << setw(25) << left << "ActivityCoefficient" << endl;
	
	BOOST_FOREACH(const string& s, species)
	{
		out << setw(25) << left << s;
		out << setw(25) << left << n[multiphase[s]];
		out << setw(25) << left << a[multiphase[s]];
		out << setw(25) << left << g[multiphase[s]] << endl;
	}
}

void GeoState::OutputReactionsState(ostream& out) const
{
	const ReactionSystem& eReactions = geosystem.GetEquilibriumReactions();
	const ReactionSystem& kReactions = geosystem.GetKineticReactions();
	
	auto outReaction = [&](const Reaction& reaction)
		{
			const double K = reaction.EquilibriumConst(T, P);
			const double Q = reaction.ReactionQuotient(a);
			const double rate = reaction.KineticRate(T, P, n, a);
			
			out << "Reaction: " << reaction << endl;
			out << "      log(K) = " << log(K) << endl;
			out << "      log(Q) = " << log(Q) << endl;
			out << "    log(Q/K) = " << log(Q/K) << endl;
			out << "        Rate = " << rate << " mol/s" << endl << endl;
		};
	
	if(eReactions.GetNumReactions() > 0)
	{
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		out << ">>> Equilibrium-Controlled Reactions <<<" << endl;
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		
		BOOST_FOREACH(auto reaction, eReactions.GetReactions())
			outReaction(reaction);
	}
	
	if(kReactions.GetNumReactions() > 0)
	{
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		out << ">>> Kinetically-Controlled Reactions <<<" << endl;
		out << ">>>>>>>>>>>>>>>>>>>><<<<<<<<<<<<<<<<<<<<" << endl;
		
		BOOST_FOREACH(auto reaction, kReactions.GetReactions())
			outReaction(reaction);
	}
}
	
ostream& operator<<(ostream& out, const GeoState& geostate)
{
	geostate.OutputMultiphaseState(out);
	
	out << endl;
	
	geostate.OutputReactionsState(out);
	
	return out;
}
